from flask import Flask, render_template, request, redirect, url_for, session
import torch
from models.detector_cnnlstm import CNNLSTMModel
import models.generator as Gen
from data_set import MyDataSet
from escape_rate import escape_rate
from torch.utils.data import DataLoader
import time

app = Flask(__name__)
app.secret_key = 'traffic'  # 设置一个秘密键，用于安全地签名session
@app.route('/index')
def index():
    """首页路由，返回index.html"""
    return render_template('index.html')

@app.route('/introduction')
def introduction():
    """介绍页面路由，返回introduction.html"""
    return render_template('introduction.html')

@app.route('/perturbation', methods=['GET', 'POST'])
def perturbation():
    if request.method == 'POST':
        # 获取用户输入的超参数
        perturbation_strength = request.form.get('perturbation_strength', type=float)
        alpha = request.form.get('alpha', type=float)
        beta = request.form.get('beta', type=float)
        zeta = request.form.get('zeta', type=float)
        mu = request.form.get('mu', type=float)
        epsilon = request.form.get('epsilon', type=float)

        # 打印获取的超参数值
        print(f"收到超参数设置："
              f"\n  扰动强度 (ξ): {perturbation_strength}"
              f"\n  对抗损失权重 (α): {alpha}"
              f"\n  生成器学习率 (β): {beta}"
              f"\n  判别器学习率 (ζ): {zeta}"
              f"\n  动量因子 (μ): {mu}"
              f"\n  集成模型学习率 (ε): {epsilon}")

        # 将超参数存储到 session 中
        session['perturbation_strength'] = perturbation_strength
        session['alpha'] = alpha
        session['beta'] = beta
        session['zeta'] = zeta
        session['mu'] = mu
        session['epsilon'] = epsilon

        # 重定向到 /generate 页面
        return redirect(url_for('generate'))
    return render_template('perturbation.html')

@app.route('/generate')
def generate():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    generator_model = Gen.Generator(input_dim=118, output_dim=118, feature_ranges={i: (-float("inf"), float("inf")) for i in range(118)}).to(device)
    generator_model.load_state_dict(torch.load("./models/generator.pth"))
    generator_model.eval()

    test_data = MyDataSet('./data/train_subset.csv')
    test_dataloader = DataLoader(test_data, batch_size=128, shuffle=True)

    # 从 session 中获取参数
    perturbation_strength = session.get('perturbation_strength', 0.8)
    alpha = session.get('alpha', 0.1)
    beta = session.get('beta', 0.03)
    zeta = session.get('zeta', 0.002)
    mu = session.get('mu', 0.6)
    epsilon = session.get('epsilon', 0.0002)

    # 打印获取的参数值（格式1）
    print(f"perturbation_strength={perturbation_strength}")
    print(f"alpha={alpha}")
    print(f"beta={beta}")
    print(f"zeta={zeta}")
    print(f"mu={mu}")
    print(f"epsilon={epsilon}")

    fake_samples = []

    with torch.no_grad():
        for real_data, real_labels in test_dataloader:  # 假设数据加载器返回 (real_data, real_labels)
            real_data = real_data.to(device)
            real_labels = real_labels.to(device)
            a,fake_sample = generator_model.generate_adv_samples(real_data, x_label=real_labels, ξ=perturbation_strength)
            fake_samples.append(fake_sample.cpu().numpy())  # 将生成的对抗样本转为 NumPy 数组并存储
            break  # 只取第一批数据

    fake_samples = fake_samples[0]  # 提取第一批对抗样本

    # 渲染 generate.html 模板，并传递超参数和生成的对抗样本
    return render_template('generate.html',
                           perturbation_strength=perturbation_strength,
                           alpha=alpha,
                           beta=beta,
                           zeta=zeta,
                           mu=mu,
                           epsilon=epsilon,
                           fake_samples=fake_samples)


@app.route('/detect')
def detect():
    start_time = time.time()  # 开始时间

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    trained_model = CNNLSTMModel().to(device)
    trained_model.load_state_dict(torch.load("./models/cnn_lstm_model.pth"))
    trained_model.eval()

    generator_model = Gen.Generator(input_dim=118, output_dim=118, feature_ranges={i: (-float("inf"), float("inf")) for i in range(118)}).to(device)
    generator_model.load_state_dict(torch.load("./models/generator.pth"))
    generator_model.eval()

    test_data = MyDataSet('./data/train_subset.csv')
    test_dataloader = DataLoader(test_data, batch_size=128, shuffle=True)

    # 从 session 中获取参数
    perturbation_strength = session.get('perturbation_strength', 0.8)
    alpha = session.get('alpha', 0.1)
    beta = session.get('beta', 0.03)
    zeta = session.get('zeta', 0.002)
    mu = session.get('mu', 0.6)
    epsilon = session.get('epsilon', 0.0002)

    # 打印获取的参数值
    print(f"在 /generate 页面获取参数：")
    print(f"  扰动强度 (ξ): {perturbation_strength}")
    print(f"  对抗损失权重 (α): {alpha}")
    print(f"  生成器学习率 (β): {beta}")
    print(f"  判别器学习率 (ζ): {zeta}")
    print(f"  动量因子 (μ): {mu}")
    print(f"  集成模型学习率 (ε): {epsilon}")

    escape_rate_val = escape_rate(model=trained_model, test_dataloader=test_dataloader, generator_model=generator_model, ξ=perturbation_strength)*7000

    end_time = time.time()  # 结束时间
    detection_time = end_time - start_time  # 计算检测时间

    # 将逃逸率和检测时间传递给前端模板
    return render_template('detect.html',
                           escape_rate=escape_rate_val,
                           perturbation_strength=perturbation_strength,
                           alpha=alpha,
                           beta=beta,
                           zeta=zeta,
                           mu=mu,
                           epsilon=epsilon,detection_time=detection_time)

if __name__ == '__main__':
    app.run(debug=True)